#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LangGraph API 集成示例
这个版本专门处理API集成，避免了复杂的参数问题
"""

import os
from typing import Dict, Any, List
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
import datetime
import re

# 尝试加载环境变量
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

# 定义工具函数
@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息"""
    weather_data = {
        "北京": "晴天，温度25°C，湿度60%，风速3m/s",
        "上海": "多云，温度22°C，湿度70%，风速2m/s", 
        "广州": "雨天，温度28°C，湿度85%，风速1m/s",
        "深圳": "晴天，温度30°C，湿度65%，风速2m/s",
        "杭州": "阴天，温度20°C，湿度75%，风速1m/s",
        "成都": "多云，温度24°C，湿度68%，风速2m/s",
        "西安": "晴天，温度26°C，湿度55%，风速3m/s",
        "武汉": "多云，温度23°C，湿度72%，风速2m/s",
        "重庆": "雨天，温度27°C，湿度80%，风速1m/s",
        "天津": "晴天，温度24°C，湿度58%，风速2m/s"
    }
    return weather_data.get(city, f"{city}的天气：晴天，温度适宜，湿度正常")

@tool
def get_time_info() -> str:
    """获取当前时间信息"""
    now = datetime.datetime.now()
    weekday = ["星期一", "星期二", "星期三", "星期四", "星期五", "星期六", "星期日"][now.weekday()]
    return f"当前时间：{now.strftime('%Y年%m月%d日')} {weekday} {now.strftime('%H:%M:%S')}"

@tool
def calculate_math(expression: str) -> str:
    """计算数学表达式"""
    try:
        allowed_chars = set('0123456789+-*/.() ')
        if not all(c in allowed_chars for c in expression):
            return "错误：只支持基本的数学运算符"
        
        if 'import' in expression or 'exec' in expression:
            return "错误：不支持的操作"
        
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算错误：{str(e)}"

@tool
def get_city_info(city: str) -> str:
    """获取城市基本信息"""
    city_info = {
        "北京": "中国首都，人口约2200万，著名景点：故宫、长城、天坛",
        "上海": "中国经济中心，人口约2400万，著名景点：外滩、东方明珠、豫园",
        "广州": "广东省省会，人口约1500万，著名景点：广州塔、陈家祠、白云山",
        "深圳": "经济特区，人口约1300万，著名景点：世界之窗、深圳湾公园、大梅沙",
        "杭州": "浙江省省会，人口约1200万，著名景点：西湖、灵隐寺、钱塘江",
        "成都": "四川省省会，人口约2100万，著名景点：宽窄巷子、大熊猫基地、武侯祠"
    }
    return city_info.get(city, f"抱歉，我还没有{city}的详细信息")

def check_api_keys():
    """检查API密钥"""
    openai_key = os.getenv("OPENAI_API_KEY")
    anthropic_key = os.getenv("ANTHROPIC_API_KEY")
    
    if openai_key and openai_key.startswith("sk-"):
        return "openai", openai_key
    elif anthropic_key:
        return "anthropic", anthropic_key
    else:
        return "none", None

def create_simple_agent():
    """创建一个简单的Agent，使用LangGraph的基本功能"""
    provider, api_key = check_api_keys()
    
    if provider == "openai":
        try:
            from langchain_openai import ChatOpenAI
            from langgraph.prebuilt import create_react_agent
            
            # 创建模型
            model = ChatOpenAI(
                model="gpt-3.5-turbo",
                temperature=0.7
            )
            
            # 创建Agent - 使用最基本的参数
            agent = create_react_agent(
                model=model,
                tools=[get_weather, get_time_info, calculate_math, get_city_info]
            )
            
            return agent, "openai"
            
        except Exception as e:
            print(f"OpenAI Agent创建失败：{str(e)}")
            return None, "failed"
    
    elif provider == "anthropic":
        try:
            from langchain_anthropic import ChatAnthropic
            from langgraph.prebuilt import create_react_agent
            
            # 创建模型
            model = ChatAnthropic(
                model="claude-3-sonnet-20240229",
                temperature=0.7
            )
            
            # 创建Agent
            agent = create_react_agent(
                model=model,
                tools=[get_weather, get_time_info, calculate_math, get_city_info]
            )
            
            return agent, "anthropic"
            
        except Exception as e:
            print(f"Anthropic Agent创建失败：{str(e)}")
            return None, "failed"
    
    return None, "none"

class LocalAgent:
    """本地Agent作为后备方案"""
    
    def __init__(self):
        self.tools = {
            "get_weather": get_weather,
            "get_time_info": get_time_info,
            "calculate_math": calculate_math,
            "get_city_info": get_city_info
        }
    
    def invoke(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """模拟LangGraph Agent的invoke方法"""
        messages = input_data.get("messages", [])
        if not messages:
            return {"messages": [AIMessage(content="请提供一个问题。")]}
        
        # 获取用户消息
        user_message = messages[-1].content if messages else ""
        
        # 简单的意图识别
        response = self._process_user_input(user_message)
        
        # 返回格式化的结果
        return {
            "messages": messages + [AIMessage(content=response)]
        }
    
    def _process_user_input(self, user_input: str) -> str:
        """处理用户输入"""
        user_input_lower = user_input.lower()
        
        # 天气询问
        if any(word in user_input_lower for word in ["天气", "weather", "温度"]):
            cities = ["北京", "上海", "广州", "深圳", "杭州", "成都", "西安", "武汉", "重庆", "天津"]
            for city in cities:
                if city in user_input:
                    result = get_weather.invoke({"city": city})
                    return f"🌤️ {city}的天气：{result}"
            result = get_weather.invoke({"city": "北京"})
            return f"🌤️ 北京的天气：{result}"
        
        # 时间询问
        elif any(word in user_input_lower for word in ["时间", "几点", "现在", "日期"]):
            result = get_time_info.invoke({})
            return f"⏰ {result}"
        
        # 数学计算
        elif any(word in user_input_lower for word in ["计算", "算"]) or any(char in user_input for char in "+-*/"):
            math_pattern = r'[\d+\-*/\.\(\)\s]+'
            match = re.search(math_pattern, user_input)
            if match:
                expression = match.group().strip()
                result = calculate_math.invoke({"expression": expression})
                return f"🔢 {result}"
        
        # 城市信息
        elif any(word in user_input_lower for word in ["城市", "介绍", "信息"]):
            cities = ["北京", "上海", "广州", "深圳", "杭州", "成都"]
            for city in cities:
                if city in user_input:
                    result = get_city_info.invoke({"city": city})
                    return f"🏙️ {city}信息：{result}"
        
        return ("🤖 我是一个智能助手，可以帮您：\n"
                "1. 查询天气：'北京天气怎么样？'\n"
                "2. 查询时间：'现在几点？'\n"
                "3. 数学计算：'计算 2+3*4'\n"
                "4. 城市信息：'介绍一下上海'")

def main():
    """主函数"""
    print("=" * 60)
    print("🤖 LangGraph API 集成示例")
    print("=" * 60)
    
    # 检查API可用性
    provider, api_key = check_api_keys()
    
    if provider != "none":
        print(f"✅ 检测到 {provider.upper()} API密钥")
        print("🔄 尝试创建API Agent...")
        
        agent, result = create_simple_agent()
        
        if agent and result in ["openai", "anthropic"]:
            print(f"🚀 成功创建 {result.upper()} Agent!")
            agent_type = "api"
        else:
            print("❌ API Agent创建失败，使用本地Agent")
            agent = LocalAgent()
            agent_type = "local"
    else:
        print("ℹ️  未检测到API密钥，使用本地Agent")
        print("💡 要使用API功能，请设置环境变量：")
        print("   export OPENAI_API_KEY=your-key")
        print("   或在.env文件中设置")
        agent = LocalAgent()
        agent_type = "local"
    
    print("\n支持的功能：")
    print("1. 天气查询：'北京天气怎么样？'")
    print("2. 时间查询：'现在几点？'")
    print("3. 数学计算：'计算 15 + 25 * 2'")
    print("4. 城市信息：'介绍一下上海'")
    print("5. 输入 'quit' 或 'exit' 退出")
    print("=" * 60)
    
    # 演示示例
    print("\n📝 演示示例：")
    demo_questions = [
        "上海天气怎么样？",
        "现在几点了？",
        "计算 12 * 8 + 5",
        "介绍一下杭州"
    ]
    
    for question in demo_questions:
        print(f"\n问题：{question}")
        
        try:
            response = agent.invoke({
                "messages": [HumanMessage(content=question)]
            })
            
            # 提取回答
            if "messages" in response and response["messages"]:
                last_message = response["messages"][-1]
                answer = last_message.content if hasattr(last_message, 'content') else str(last_message)
                print(f"回答：{answer}")
            else:
                print("回答：处理出错")
                
        except Exception as e:
            print(f"处理错误：{str(e)}")
        
        print("-" * 40)
    
    # 交互式对话
    print("\n🎯 现在您可以开始提问了！")
    
    while True:
        try:
            user_input = input("\n您的问题：").strip()
            
            if user_input.lower() in ['quit', 'exit', '退出', 'bye', 'q']:
                print("👋 再见！感谢使用LangGraph助手！")
                break
            
            if not user_input:
                print("请输入有效的问题。")
                continue
            
            response = agent.invoke({
                "messages": [HumanMessage(content=user_input)]
            })
            
            # 提取回答
            if "messages" in response and response["messages"]:
                last_message = response["messages"][-1]
                answer = last_message.content if hasattr(last_message, 'content') else str(last_message)
                print(f"🤖 助手：{answer}")
            else:
                print("🤖 助手：处理出错")
                
        except KeyboardInterrupt:
            print("\n\n👋 程序已终止。再见！")
            break
        except Exception as e:
            print(f"❌ 发生错误：{str(e)}")
            print("请重试...")

if __name__ == "__main__":
    main() 